#include "servbuffer.h"
#include "lspmessage.pb.h"
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>
#include <arpa/inet.h>
#include <ctime>
#include <iostream>
#define UDP_PAYLOAD_LEN 65507
#define EPOCH_LEN 3
#define EPOCH_EXP 3

using namespace std;

void * serv_io_thread_func(void * arg) {
	printf("Starting I/O Thread\n");
	ServerBuffer * serv = (ServerBuffer*)arg;
	
	// Time data
	time_t epoch_timer = time(NULL);

	while(serv->running) {
		// Polling data
		fd_set fds;
		FD_ZERO(&fds);
		FD_SET(serv->sock, &fds);
		struct timeval timeout;
		timeout.tv_usec = 100;
		timeout.tv_sec = 0;
		
		// wait until socket is ready
		int ready = select(serv->sock+1, &fds, 0, 0, &timeout);
		if(ready > 0) {
			// Read and process packet
			//printf("Starting ready loop\n");
			struct sockaddr_in recv_addr;
			socklen_t recv_addr_len = sizeof(struct sockaddr_in);
			char buf[UDP_PAYLOAD_LEN];
			memset(buf, 0, UDP_PAYLOAD_LEN);
			//printf("After clearing buffer\n");

			int bytes = recvfrom(serv->sock, buf, UDP_PAYLOAD_LEN, 0,
					(struct sockaddr*)&recv_addr, &recv_addr_len);
			if(bytes == -1)
				serv->print_error(errno, "Failed to read from socket");
			//else if(bytes < 8)
			//	serv->print_error(-1, "Invalid server header");
			//printf("received packet\n");

			Packet dummy((const uint8_t*)buf, bytes);
			//printf("Created dummy packet\n");
			uint32_t conn_id = dummy.get_conn_id();
			uint32_t seq = dummy.get_seq_num();
			const uint8_t * pld = dummy.get_payload();
			int pld_len = strlen((const char*)pld);
			//printf("Conn id: %d. Sequence number: %d. Payload: %s. len: %d.\n", conn_id, seq, dummy.get_payload(), pld_len);
			
			int cli_id = 0;
			if(serv->connected_clients.size() > 0)
				cli_id = serv->connected_clients.front()->conn_id;

			list<ClientData*>::iterator client = serv->connected_clients.begin();
			while( client != serv->connected_clients.end() &&
				!(*client)->compare(&recv_addr, conn_id) )
				client++;
			
			if(conn_id == 0) {
				//printf("Case 1\n");
				if(client == serv->connected_clients.end()) {
					// New client connection
					ClientData* new_client = new ClientData();

					new_client->addr = recv_addr;

					new_client->conn_id = serv->get_next_id();
					new_client->snd_seq = 1;
					new_client->rcv_seq = 1;
					new_client->last_io_event = time(NULL);
					pthread_mutex_init(&new_client->in_lock, NULL);
					pthread_mutex_init(&new_client->out_lock, NULL);
					client = serv->insert_client(new_client);
					printf("Client connected at %s:%d\n", inet_ntoa(new_client->addr.sin_addr),
										ntohs(new_client->addr.sin_port));
				}

				// Create setup response packet
				Packet setup((*client)->conn_id, 0, 0, 0);
				
				// Send it off
				string setup_str;
				setup.msg.SerializeToString(&setup_str);
				bytes = sendto(serv->sock, setup_str.c_str(), setup_str.size(), 0,
						(struct sockaddr*)&recv_addr, recv_addr_len);
				if(bytes == -1)
					serv->print_error(errno, "Failed to send packet to client");
				(*client)->last_io_event = time(NULL);
				(*client)->epoch_counter = 0;
			}
			else if(pld_len == 0 && client != serv->connected_clients.end()) {
				//printf("Case 2\n");
				// Client acknowledgement
				list<Packet>::iterator segment = (*client)->out_buf.begin();
				while(segment != (*client)->out_buf.end() &&
						segment->get_seq_num() != seq)
					segment++;
				if(segment != (*client)->out_buf.end()) {
					// Packet is acknowledged
					pthread_mutex_lock(&((*client)->out_lock));
					(*client)->out_buf.erase(segment, (*client)->out_buf.end());
					pthread_mutex_unlock(&((*client)->out_lock));
					(*client)->snd_seq = seq + 1;

					(*client)->last_io_event = time(NULL);
					(*client)->epoch_counter = 0;
				}

				if(segment == (*client)->out_buf.end())
					segment--;

				// Send first unacknowledged packet
				if((*client)->out_buf.size() > 0) {
					Packet p = *segment;
					string p_str;
					p.msg.SerializeToString(&p_str);
					bytes = sendto(serv->sock, p_str.c_str(), p_str.size(), 0,
							(struct sockaddr*)&((*client)->addr),
							sizeof(struct sockaddr_in));
					if(bytes == -1)
						serv->print_error(errno, "Failed to send packet to client");
				}

				// Liveness
				if(seq == (*client)->snd_seq-1) {
					(*client)->last_io_event = time(NULL);
					(*client)->epoch_counter = 0;
				}
			}
			else if(pld_len > 0 && client != serv->connected_clients.end()) {
				//printf("Case 3\n");
				// Client data segment
				if(seq == (*client)->rcv_seq) {
					// Packet is expected
					Packet data((uint8_t*)buf, bytes);
					pthread_mutex_lock(&((*client)->in_lock));
					(*client)->in_buf.push_front(data);
					pthread_mutex_unlock(&((*client)->in_lock));
					(*client)->rcv_seq += 1;
				}
				//printf("Case 3 after if statement\n");

				// Ack the last acknowledged packet
				Packet ack(conn_id, (*client)->rcv_seq-1, 0, 0);
				//printf("Created ack packet\n");
				string ack_str;
				ack.msg.SerializeToString(&ack_str);
				//printf("Serialized ack string\n");
 				bytes = sendto(serv->sock, ack_str.c_str(), ack_str.size(), 0,
						(struct sockaddr*)&((*client)->addr),
						sizeof(struct sockaddr_in));
						//printf("sent bytes\n");
				if(bytes == -1)
					serv->print_error(errno, "Failed to send packet to client");
			}
			else {
				printf("Bad request\n");
			}
			(*client)->last_io_event = time(NULL);
			(*client)->epoch_counter = 0;
		}

		// Handle epoch events
		time_t current_time = time(NULL);
		if(current_time - epoch_timer >= serv->epoch) {

			for(list<ClientData*>::iterator clients = serv->connected_clients.begin();
					clients != serv->connected_clients.end();
					clients++) {

				// Send unacknowledged packet
				if((*clients)->out_buf.size() > 0) {
					Packet p = (*clients)->out_buf.back();
					string p_str;
					p.msg.SerializeToString(&p_str);
					int bytes = sendto(serv->sock, p_str.c_str(), p_str.size(), 0,
							(struct sockaddr*)&((*clients)->addr),
							sizeof(struct sockaddr_in));
					if(bytes == -1)
						serv->print_error(errno, "Failed to send packet to client");
				}

				// Acknowledge latest received packet
				Packet ack((*clients)->conn_id, (*clients)->rcv_seq-1, 0, 0);
				string ack_str;
				ack.msg.SerializeToString(&ack_str);
				int bytes = sendto(serv->sock, ack_str.c_str(), ack_str.size(), 0,
							(struct sockaddr*)&((*clients)->addr),
							sizeof(struct sockaddr_in));
				if(bytes == -1)
					serv->print_error(errno, "Failed to send packet to client");

				// Resend connection setups
				if((*clients)->rcv_seq == 1 && (*clients)->snd_seq == 1) {
					Packet setup((*clients)->conn_id, 0, 0, 0);
					string setup_str;
					setup.msg.SerializeToString(&setup_str);
					bytes = sendto(serv->sock, setup_str.c_str(), setup_str.size(), 0,
								(struct sockaddr*)&((*clients)->addr),
								sizeof(struct sockaddr_in));
					if(bytes == -1)
						serv->print_error(errno, "Failed to send packet to client");
				}
			}
			epoch_timer = time(NULL);
		}

		// Handle timeouts
		for(list<ClientData*>::iterator iter = serv->connected_clients.begin();
				iter != serv->connected_clients.end();
				iter++) {
			if(current_time - (*iter)->last_io_event >= serv->epoch) {
				(*iter)->epoch_counter++;
				(*iter)->last_io_event = time(NULL);
			}
			if((*iter)->epoch_counter == 3) {
				// Connection has timed out; delete client
//				printf("Erasing client: %s:%d\n", inet_ntoa((*iter)->addr.sin_addr), ntohs((*iter)->addr.sin_port));
				iter = serv->connected_clients.erase(iter);
			}
		}
	}

	printf("Shutting down server.\n");
}

ServerBuffer::ServerBuffer(uint16_t port, int epoch_len, int n) {
	server.sin_addr.s_addr = INADDR_ANY;
	server.sin_port = htons(port);
	server.sin_family = AF_INET;
	state = 0;
	epoch = epoch_len;
	numepochs = n;
	pthread_mutex_init(&client_lock, NULL);

	sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if(sock == -1) {
		state = errno;
		errmsg = "Failed to create socket";
		return;
	}
	int result = bind(sock, (struct sockaddr*)&server, sizeof(server));
	if(result == -1) {
		state = errno;
		errmsg = "Failed to bind to port";
		return;
	}
	running = true;
	
	result = pthread_create(&io_thread, 0, serv_io_thread_func, this);
	if(result == -1) {
		state = errno;
		errmsg = "Failed to spawn I/O thread";
		return;
	}
}

void ServerBuffer::finish() {
	running = false;
	pthread_join(io_thread, NULL);
	close(sock);
}

bool ClientData::compare(struct sockaddr_in * a, uint32_t c_id) {
	if(c_id == 0) {
		return (addr.sin_addr.s_addr == a->sin_addr.s_addr
			&& addr.sin_port == a->sin_port);
	}
	else {
		return (addr.sin_addr.s_addr == a->sin_addr.s_addr
                        && addr.sin_port == a->sin_port
			&& conn_id == c_id);
	}
}

void ServerBuffer::print_error(int err, const char * msg) {
	printf("%s: %s.\n", msg, strerror(errno));
	this->state = 0;
}

uint32_t ServerBuffer::get_next_id() {
	list<ClientData*>::iterator iter = this->connected_clients.begin();
	uint32_t id = 1;
	while(iter != this->connected_clients.end() &&
			id == (*iter)->conn_id) {
		iter++;
		id++;
	}
	return id;
}

list<ClientData*>::iterator ServerBuffer::insert_client(ClientData* new_client) {
	list<ClientData*>::iterator iter = this->connected_clients.begin();
	while(iter != this->connected_clients.end() &&
			(*iter)->conn_id < new_client->conn_id)
		iter++;
	return this->connected_clients.insert(iter, new_client);
}

int ServerBuffer::read(uint8_t * data, uint32_t* conn_id) {
	list<ClientData*>::iterator client = connected_clients.begin();
	while(client != connected_clients.end() &&
		(*client)->in_buf.size() == 0) {
		client++;
	}

	if(client != connected_clients.end()) {
		*conn_id = (*client)->conn_id;
		Packet p = (*client)->in_buf.front();
		memcpy(data, (const char*)p.get_payload(), strlen((const char*)p.get_payload()));
		pthread_mutex_lock(&((*client)->in_lock));
		(*client)->in_buf.pop_front();
		pthread_mutex_unlock(&((*client)->in_lock));
		return strlen((const char*)p.get_payload());
	}
	else {
		return 0;
	}
}

int ServerBuffer::write(uint32_t conn_id, uint8_t * data, uint32_t len) {
	list<ClientData*>::iterator client = connected_clients.begin();
	while(client != connected_clients.end() &&
                (*client)->conn_id != conn_id)  
                client++;
	
	if(client != connected_clients.end()) {
		pthread_mutex_lock(&((*client)->out_lock));
		uint32_t seq = (*client)->snd_seq;
		if((*client)->out_buf.size() > 0)
			seq = (*client)->out_buf.front().get_seq_num()+1;
		Packet p(conn_id, seq, data, len);
		string p_str;
		p.msg.SerializeToString(&p_str);
		(*client)->out_buf.push_front(p);
		pthread_mutex_unlock(&((*client)->out_lock));
		sendto(sock, p_str.c_str(), p_str.size(), 0, (struct sockaddr*)&((*client)->addr), sizeof(struct sockaddr_in));
		return len;
	}
	else {
		return -1;
	}
}

